import time
import numpy as np
import cpu_extension

def run_numpy_matmul(A, B):
    start = time.time()
    C = np.dot(A, B)
    return C, time.time() - start

def run_cpp_matmul(A, B):
    start = time.time()
    C = cpu_extension.matmul(A, B)
    return C, time.time() - start

def benchmark(size):
    A = np.random.rand(size, size).astype(np.float32)
    B = np.random.rand(size, size).astype(np.float32)

    print(f"\n🔹 測試矩陣大小: {size} x {size}")

    C_np, time_np = run_numpy_matmul(A, B)
    print(f"  NumPy dot 耗時：{time_np:.4f} 秒")

    C_cpp, time_cpp = run_cpp_matmul(A, B)
    print(f"  C++ matmul 耗時：{time_cpp:.4f} 秒")

    diff = np.max(np.abs(C_np - C_cpp))
    print(f"  結果最大誤差：{diff:.6f}")

if __name__ == "__main__":
    for size in [256, 512, 1024]:
        benchmark(size)